package com.starhub.integration.sdk.shangtang.model;

import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.model.output.TokenUsage;
import okhttp3.*;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.security.SecureRandom;
import java.security.cert.X509Certificate;
import java.util.List;
import java.util.concurrent.TimeUnit;
import dev.langchain4j.model.output.FinishReason;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;

/**
 * 商汤模型适配器（非流式）
 * 无需认证，直接调用本地商汤API
 */
public class ShangtangChatModel implements ChatLanguageModel {

    private static final Logger log = LoggerFactory.getLogger(ShangtangChatModel.class);

    private final String apiUrl;
    private final OkHttpClient httpClient;
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final String modelName;
    private final double temperature;
    private final int maxTokens;

    public ShangtangChatModel(String apiUrl, String modelName, double temperature, int maxTokens) {
        this.apiUrl = apiUrl;
        this.modelName = modelName;
        this.temperature = temperature;
        this.maxTokens = maxTokens;

        // 修改为创建绕过SSL验证的HTTP客户端
        this.httpClient = createUnsafeOkHttpClient();
    }

    // 创建绕过SSL验证的HTTP客户端
    private OkHttpClient createUnsafeOkHttpClient() {
        try {
            // 1. 创建信任所有证书的TrustManager
            TrustManager[] trustAllCerts = new TrustManager[]{
                    new X509TrustManager() {
                        public void checkClientTrusted(X509Certificate[] chain, String authType) {}
                        public void checkServerTrusted(X509Certificate[] chain, String authType) {}
                        public X509Certificate[] getAcceptedIssuers() {
                            return new X509Certificate[]{};
                        }
                    }
            };

            // 2. 创建SSLContext
            SSLContext sslContext = SSLContext.getInstance("SSL");
            sslContext.init(null, trustAllCerts, new SecureRandom());

            // 3. 创建OkHttpClient
            return new OkHttpClient.Builder()
                    .sslSocketFactory(sslContext.getSocketFactory(), (X509TrustManager) trustAllCerts[0])
                    .hostnameVerifier((hostname, session) -> true) // 忽略主机名验证
                    .connectTimeout(60, TimeUnit.SECONDS)
                    .readTimeout(60, TimeUnit.SECONDS)
                    .build();
        } catch (Exception e) {
            throw new RuntimeException("创建HTTP客户端失败", e);
        }
    }

    @Override
    public Response<AiMessage> generate(List<ChatMessage> messages) {
        try {
            // 1. 构建商汤API请求
            ShangtangRequest request = new ShangtangRequest(
                    modelName,
                    convertMessages(messages),
                    false, // 非流式
                    temperature,
                    maxTokens
            );
            log.info("ShangtangChatModel request>>>>>>>>>>>>>>>"+objectMapper.writeValueAsString(request));
            // 2. 发送HTTP请求
            Request httpRequest = new Request.Builder()
                    .url(apiUrl)
                    .post(RequestBody.create(
                            objectMapper.writeValueAsString(request),
                            MediaType.get("application/json")
                    ))
                    .build();
            log.info("ShangtangChatModel apiUrl>>>>>>>>>>>>>>>"+ apiUrl);
            // 3. 执行请求并解析响应
            try (okhttp3.Response response = httpClient.newCall(httpRequest).execute()) {
                log.info("ShangtangChatModel response>>>>>>>>>>>>>>>"+ JSONObject.toJSONString(response));
                if (!response.isSuccessful()) {
                    throw new RuntimeException("商汤API请求失败: " + response.code());
                }
                ShangtangResponse shangtangResponse = objectMapper.readValue(
                        response.body().string(),
                        ShangtangResponse.class
                );
                // 4. 转换为LangChain4j响应格式
                return convertResponse(shangtangResponse);
            }
        } catch (Exception e) {
            log.error("ShangtangResponse",e);
            throw new RuntimeException("调用商汤模型失败", e);
        }
    }

    // 转换消息格式
    private List<ShangtangMessage> convertMessages(List<ChatMessage> messages) {
        return messages.stream()
                .map(msg -> new ShangtangMessage(
                        msg.type().name().toLowerCase(),
                        msg.text()
                ))
                .collect(java.util.stream.Collectors.toList());
    }

    // 转换响应格式
    private Response<AiMessage> convertResponse(ShangtangResponse response) {
        if (response.getChoices() == null || response.getChoices().isEmpty()) {
            throw new RuntimeException("商汤API返回空响应");
        }

        ShangtangResponse.Choice choice = response.getChoices().get(0);
        String content = choice.getMessage().getContent();

        TokenUsage tokenUsage = null;
        if (response.getUsage() != null) {
            tokenUsage = new TokenUsage(
                    response.getUsage().getPrompt_tokens(),
                    response.getUsage().getCompletion_tokens(),
                    response.getUsage().getTotal_tokens()
            );
        }

        return Response.from(
                new AiMessage(content),
                tokenUsage,
                FinishReason.STOP
        );
    }

    // 商汤API请求结构
    private static class ShangtangRequest {
        private final String model;
        private final List<ShangtangMessage> messages;
        private final boolean stream;
        private final double temperature;
        private final int max_tokens;

        public ShangtangRequest(String model, List<ShangtangMessage> messages,
                                boolean stream, double temperature, int max_tokens) {
            this.model = model;
            this.messages = messages;
            this.stream = stream;
            this.temperature = temperature;
            this.max_tokens = max_tokens;
        }

        // Getters
        public String getModel() { return model; }
        public List<ShangtangMessage> getMessages() { return messages; }
        public boolean isStream() { return stream; }
        public double getTemperature() { return temperature; }
        public int getMax_tokens() { return max_tokens; }
    }

    // 消息结构
    private static class ShangtangMessage {
        private final String role;
        private final String content;

        public ShangtangMessage(String role, String content) {
            this.role = role;
            this.content = content;
        }

        // Getters
        public String getRole() { return role; }
        public String getContent() { return content; }
    }

    // 商汤API响应结构
    @JsonIgnoreProperties(ignoreUnknown = true)
    // 商汤API响应结构
    private static class ShangtangResponse {
        private String id;
        private String object;
        private long created;
        private String model;
        private List<Choice> choices;
        private Usage usage;

        // Getters and setters
        public String getId() { return id; }
        public void setId(String id) { this.id = id; }

        public String getObject() { return object; }
        public void setObject(String object) { this.object = object; }

        public long getCreated() { return created; }
        public void setCreated(long created) { this.created = created; }

        public String getModel() { return model; }
        public void setModel(String model) { this.model = model; }

        public List<Choice> getChoices() { return choices; }
        public void setChoices(List<Choice> choices) { this.choices = choices; }

        public Usage getUsage() { return usage; }
        public void setUsage(Usage usage) { this.usage = usage; }

        public static class Choice {
            private String finish_reason;
            private int index;
            private double logprobs;
            private Message message;

            // Getters and setters
            public String getFinish_reason() { return finish_reason; }
            public void setFinish_reason(String finish_reason) { this.finish_reason = finish_reason; }

            public int getIndex() { return index; }
            public void setIndex(int index) { this.index = index; }

            public double getLogprobs() { return logprobs; }
            public void setLogprobs(double logprobs) { this.logprobs = logprobs; }

            public Message getMessage() { return message; }
            public void setMessage(Message message) { this.message = message; }
        }

        public static class Message {
            private String content;
            private String role;

            // Getters and setters
            public String getContent() { return content; }
            public void setContent(String content) { this.content = content; }

            public String getRole() { return role; }
            public void setRole(String role) { this.role = role; }
        }

        public static class Usage {
            private int prompt_tokens;
            private int completion_tokens;
            private int total_tokens;

            // Getters and setters
            public int getPrompt_tokens() { return prompt_tokens; }
            public void setPrompt_tokens(int prompt_tokens) { this.prompt_tokens = prompt_tokens; }

            public int getCompletion_tokens() { return completion_tokens; }
            public void setCompletion_tokens(int completion_tokens) { this.completion_tokens = completion_tokens; }

            public int getTotal_tokens() { return total_tokens; }
            public void setTotal_tokens(int total_tokens) { this.total_tokens = total_tokens; }
        }
    }
}
